import math
import os
import random

import cv2 as cv
import keras.backend as K
import numpy as np

from data_generator import generate_trimap, random_choice
from model import build_encoder_decoder, build_refinement
from utils import compute_mse_loss, compute_sad_loss
from utils import get_final_output, safe_crop, draw_str

import glob
import platform


def _WIN32():
    return str(platform.system()).upper() == 'WINDOWS'


def GetFiles(regx):
    if _WIN32():
        return [str(f).replace('\\', '/') for f in glob.glob(regx)]
    else:
        return glob.glob(regx)


def composite4(fg, bg, a, w, h):
    fg = np.array(fg, np.float32)
    bg_h, bg_w = bg.shape[:2]
    x = 0
    if bg_w > w:
        x = np.random.randint(0, bg_w - w)
    y = 0
    if bg_h > h:
        y = np.random.randint(0, bg_h - h)
    bg = np.array(bg[y:y + h, x:x + w], np.float32)
    alpha = np.zeros((h, w, 1), np.float32)
    alpha[:, :, 0] = a / 255.
    im = alpha * fg + (1 - alpha) * bg
    im = im.astype(np.uint8)
    return im, bg


def CenterCrop(im_, dst_size):
    if type(dst_size) == int:
        dst_size = [dst_size]*2
    c = 1
    if len(im_.shape)==3:
        w = im_.shape[1]
        h = im_.shape[0]
        c = im_.shape[2]
    else:
        assert len(im_.shape)==2
        w = im_.shape[1]
        h = im_.shape[0]
    x = w - dst_size[0]
    y = h - dst_size[1]
    assert x >= 0 and y >= 0
    x = x // 2
    y = y // 2
    if c > 1:
        im_new = np.zeros([*dst_size[::-1], c], dtype=np.uint8)
    else:
        im_new = np.zeros(dst_size[::-1], dtype=np.uint8)
    im_new[:,:] = im_[y:y+dst_size[1], x:x+dst_size[0]]
    return im_new


def CenterPad(im_, dst_size):
    if type(dst_size) == int:
        dst_size = [dst_size]*2
    c = 1
    if len(im_.shape)==3:
        w = im_.shape[1]
        h = im_.shape[0]
        c = im_.shape[2]
    else:
        assert len(im_.shape)==2
        w = im_.shape[1]
        h = im_.shape[0]
    x = (dst_size[0] - w)
    y = (dst_size[1] - h)
    assert x >= 0 and y >= 0
    # if x > 0 and y > 0:
    #     ratio = min(1.0*dst_size[0]/w, 1.0*dst_size[1]/h)
    #     w, h = int(w*ratio), int(h*ratio)
    #     im_ = cv.resize(im_, (w, h))
    x = (dst_size[0] - w) // 2
    y = (dst_size[1] - h) // 2
    if c > 1:
        im_new = np.zeros([*dst_size[::-1], c], dtype=np.uint8)
    else:
        im_new = np.zeros(dst_size[::-1], dtype=np.uint8)
    im_new[y:y+h, x:x+w] = im_[:,:]
    return im_new, (x, y, w, h)


if __name__ == '__main__':
    img_rows, img_cols = 192, 256
    channel = 4

    pretrained_path = '../Models/DIM/final.42-0.0398.hdf5'
    encoder_decoder = build_encoder_decoder(img_rows, img_cols, train=False)

    #exit(0)

    final = build_refinement(encoder_decoder, train=False)
    final.load_weights(pretrained_path)
    print(final.summary())

    pngs = GetFiles('./data/segment/*.PNG')

    for i in range(len(pngs)):
        png = CenterCrop(cv.imread(pngs[i], -1), [img_cols, img_rows])
        #png, roi = CenterPad(cv.imread(pngs[i], -1), img_cols)
        assert len(png.shape)==3 and png.shape[-1]==4
        _, alpha = cv.threshold(png[:, :, 3], 200, 255, cv.THRESH_BINARY) # 254
        ori = png[:,:,:3]
        fg = cv.erode(alpha, np.array([]), iterations=3) #3
        bg = 255 - cv.dilate(alpha, np.array([]), iterations=12) #12
        grey = 255 - cv.bitwise_or(fg, bg)
        tri = fg + cv.bitwise_and(grey, 128)

        # # remove false grey mask in trimaps
        # tri_tmp = np.zeros([img_rows, img_cols], np.uint8)
        # tri_tmp[roi[1]:roi[1]+roi[3], roi[0]:roi[0]+roi[2]] = tri[roi[1]:roi[1]+roi[3], roi[0]:roi[0]+roi[2]]
        # tri = tri_tmp

        cv.imshow('alpha', alpha)
        cv.imshow('origin', ori)
        cv.imshow('trimap', tri)
        cv.waitKey(0)

        x_test = np.empty((1, img_rows, img_cols, 4), dtype=np.float32)
        x_test[0, :, :, 0:3] = ori[:,:,:3] / 255.
        x_test[0, :, :, 3] = tri / 255.

        y_pred = final.predict(x_test)
        y_pred = np.reshape(y_pred, (img_rows, img_cols))
        y_pred = y_pred * 255.0

        # show the temporal result
        #res_u8 = np.array(y_pred, np.uint8)
        #cv.imshow('res_u8', res_u8)
        #cv.waitKey(0)

        y_pred = get_final_output(y_pred, tri)
        y_pred = y_pred.astype(np.uint8)
        mask = y_pred
        #_, mask = cv.threshold(mask, 128, 255, cv.THRESH_TOZERO)
        #mask = cv.resize(y_pred, (ori.shape[1], ori.shape[0]))
        #mask = cv.dilate(mask, np.array([]), iterations=1)
        #mask = cv.erode(mask, np.array([]), iterations=2)
        #mask = cv.medianBlur(mask, 3)
        merged = np.stack([ori[:, :, 0], ori[:, :, 1], ori[:, :, 2], mask], axis=-1)
        cv.imwrite(pngs[i].replace('segment', 'matting'), merged)
        cv.imshow('Matting Result', mask)
        cv.waitKey(0)

    K.clear_session()
